import pandas as pd
import numpy as np
import yfinance as yf
import os
import re
from datetime import datetime,timedelta
import warnings
import matplotlib.pyplot as plt
from fontTools.misc.plistlib import end_data
from pandas.io.common import file_path_to_url

warnings.filterwarnings('ignore')


def calculate_atr(df, window=14, high_col='最高', low_col='最低', close_col='收盘'):
    """
    计算平均真实波幅(ATR)
    参数:
        df: 包含价格数据的DataFrame
        window: ATR计算周期
        high_col: 最高价列名
        low_col: 最低价列名
        close_col: 收盘价列名
    返回:
        包含TR和ATR的DataFrame
    """
    # 复制DataFrame以避免修改原始数据
    df = df.copy()

    # 1. 计算真实波幅(TR)
    df['prev_close'] = df[close_col].shift(1)
    df['tr0'] = df[high_col] - df[low_col]  # 当日高-低
    df['tr1'] = abs(df[high_col] - df['prev_close'])  # 当日高-前日收
    df['tr2'] = abs(df[low_col] - df['prev_close'])  # 当日低-前日收
    df['tr'] = df[['tr0', 'tr1', 'tr2']].max(axis=1)

    # 2. 计算ATR (使用Wilder平滑方法)
    # 初始ATR为前window天的TR平均值
    df['atr'] = df['tr'].rolling(window=window, min_periods=1).mean()

    # 对于后续值，使用递归平均
    for i in range(window, len(df)):
        df.loc[df.index[i], 'atr'] = (df.loc[df.index[i - 1], 'atr'] * (window - 1) + df.loc[
            df.index[i], 'tr']) / window

    # 仅保留有用列
    return df[['tr', 'atr']]
def calculate_adx(df, date_col='日期', high_col='最高', low_col='最低', close_col='收盘', window=14):
    """计算平均方向指数(ADX)"""
    # 确保数据排序正确
    df = df.sort_values(date_col).reset_index(drop=True)

    # 计算+DM和-DM
    df['high_diff'] = df[high_col].diff()
    df['low_diff'] = -df[low_col].diff()
    df['plus_dm'] = np.where((df['high_diff'] > df['low_diff']) & (df['high_diff'] > 0), df['high_diff'], 0)
    df['minus_dm'] = np.where((df['low_diff'] > df['high_diff']) & (df['low_diff'] > 0), df['low_diff'], 0)

    # 计算真实波幅(TR)
    df['tr1'] = df[high_col] - df[low_col]
    df['tr2'] = abs(df[high_col] - df[close_col].shift(1))
    df['tr3'] = abs(df[low_col] - df[close_col].shift(1))
    df['tr'] = df[['tr1', 'tr2', 'tr3']].max(axis=1)

    # 平滑DM和TR
    df['smooth_plus_dm'] = df['plus_dm'].ewm(alpha=1 / window, adjust=False).mean()
    df['smooth_minus_dm'] = df['minus_dm'].ewm(alpha=1 / window, adjust=False).mean()
    df['smooth_tr'] = df['tr'].ewm(alpha=1 / window, adjust=False).mean()

    # 计算方向指标(+DI和-DI)
    df['plus_di'] = (df['smooth_plus_dm'] / df['smooth_tr']) * 100
    df['minus_di'] = (df['smooth_minus_dm'] / df['smooth_tr']) * 100

    # 计算方向差异和方向指标
    df['dx'] = (abs(df['plus_di'] - df['minus_di']) / (df['plus_di'] + df['minus_di'])) * 100

    # 计算ADX
    df['adx'] = df['dx'].ewm(alpha=1 / window, adjust=False).mean()

    # 返回最终ADX值
    return df['adx'].iloc[-1]


def calculate_continuous_down_months(df, date_col='日期', close_col='收盘'):
    """计算连续下跌月数"""
    # 创建月份列
    df['month'] = df[date_col].dt.to_period('M')

    # 按月聚合
    monthly = df.groupby('month').agg(
        first_date=pd.NamedAgg(column=date_col, aggfunc='min'),
        last_close=pd.NamedAgg(column=close_col, aggfunc='last')
    ).reset_index()

    # 计算月收益率
    monthly['prev_close'] = monthly['last_close'].shift(1)
    monthly['monthly_return'] = (monthly['last_close'] - monthly['prev_close']) / monthly['prev_close']

    # 标记下跌月份
    monthly['down'] = monthly['monthly_return'] < 0

    # 计算连续下跌月数
    monthly['streak'] = monthly['down'].groupby((~monthly['down']).cumsum()).cumsum()
    max_streak = monthly['streak'].max()

    # 如果全是上涨月，返回0
    return max_streak if not pd.isna(max_streak) else 0


def calculate_correlations(data_dict, date_col='日期', close_col='收盘'):
    """计算多个标的之间的相关性"""
    # 创建合并的数据集
    merged = pd.DataFrame()

    # 收集所有数据
    for name, df in data_dict.items():
        df = df.sort_values(date_col)
        temp = df[[date_col, close_col]].copy()
        temp.columns = [date_col, name]

        if merged.empty:
            merged = temp
        else:
            merged = pd.merge(merged, temp, on=date_col, how='outer')

    # 按日期排序并填充缺失值
    merged = merged.sort_values(date_col).fillna(method='ffill').dropna()

    # 计算对数收益率
    returns = pd.DataFrame()
    for col in merged.columns:
        if col != date_col:
            returns[col] = np.log(merged[col] / merged[col].shift(1))

    # 移除第一行NaN
    returns = returns.dropna()

    # 计算相关性矩阵
    corr_matrix = returns.corr()

    return corr_matrix


def calculate_all_metrics(file_path):
    """计算单个文件的所有指标"""
    try:
        # 尝试自动检测列名
        df_sample = pd.read_csv(file_path, nrows=1)

        # 日期列
        date_cols = [col for col in df_sample.columns if any(keyword in col.lower()
                                                             for keyword in
                                                             ['date', '时间', '日期', 'day', '交易时间'])]
        date_col = date_cols[0] if date_cols else '日期'

        # 价格列
        close_cols = [col for col in df_sample.columns if any(keyword in col.lower()
                                                              for keyword in ['close', '收盘', 'price'])]
        close_col = close_cols[0] if close_cols else '收盘'

        high_cols = [col for col in df_sample.columns if any(keyword in col.lower()
                                                             for keyword in ['high', '最高', '高'])]
        high_col = high_cols[0] if high_cols else '最高'

        low_cols = [col for col in df_sample.columns if any(keyword in col.lower()
                                                            for keyword in ['low', '最低', '低'])]
        low_col = low_cols[0] if low_cols else '最低'

        # 读取完整数据
        df = pd.read_csv(file_path, parse_dates=[date_col])

        # 检查必要列
        required_cols = [date_col, close_col]
        if not all(col in df.columns for col in required_cols):
            missing = [col for col in required_cols if col not in df.columns]
            print(f"⚠️ 文件 {os.path.basename(file_path)} 缺少必要列: {', '.join(missing)}")
            return None

        # 按日期排序
        df = df.sort_values(date_col).reset_index(drop=True)

        # 检查数据量
        if len(df) < 30:
            print(f"⚠️ 文件 {os.path.basename(file_path)} 数据量不足 ({len(df)} 行)，跳过计算")
            return None

        # 计算年化波动率
        df['日收益率'] = np.log(df[close_col] / df[close_col].shift(1))
        daily_std = df['日收益率'].std(ddof=1)
        annualized_vol = daily_std * np.sqrt(252)

        # 计算ADX (需要最高价和最低价)
        adx_value = None
        if high_col in df.columns and low_col in df.columns:
            try:
                adx_value = calculate_adx(df, date_col, high_col, low_col, close_col)
            except Exception as e:
                print(f"⚠️ 文件 {os.path.basename(file_path)} ADX计算失败: {str(e)}")

        # 计算连续下跌月数
        try:
            down_months = calculate_continuous_down_months(df, date_col, close_col)
        except Exception as e:
            print(f"⚠️ 文件 {os.path.basename(file_path)} 连续下跌月数计算失败: {str(e)}")
            down_months = 0
        # ============ 新增：ATR计算 ============
        atr_20 = None
        atr_ratio = None
        spot_price = None

        if high_col in df.columns and low_col in df.columns:
            try:
                # 计算ATR
                atr_df = calculate_atr(df, window=20,
                                       high_col=high_col,
                                       low_col=low_col,
                                       close_col=close_col)
                df = pd.concat([df, atr_df], axis=1)

                # 获取最新值
                latest = df.iloc[-1]
                spot_price = latest[close_col]  # 现货价 = 最新收盘价
                atr_20 = latest['atr']  # 20日ATR

                # 计算波动率比率
                atr_ratio = atr_20 / spot_price if spot_price > 0 else 0

            except Exception as e:
                print(f"⚠️ 文件 {os.path.basename(file_path)} ATR计算失败: {str(e)}")
        # =====================================
        return {
            '文件名称': os.path.basename(file_path),
            '数据量': len(df),
            '开始日期': df[date_col].min().strftime('%Y-%m-%d'),
            '结束日期': df[date_col].max().strftime('%Y-%m-%d'),
            '年化波动率': annualized_vol,
            'ADX(14日)': adx_value,
            '连续下跌月数': down_months,
            '20日ATR': atr_20,
            '现货价': spot_price,
            'ATR比率(%)': atr_ratio * 100 if atr_ratio is not None else None,
            '收盘价列': close_col,
            '数据框': df[[date_col, close_col]]
        }

    except Exception as e:
        print(f"❌ 处理文件 {os.path.basename(file_path)} 时出错: {str(e)}")
        return None


def analyze_directory(directory_path):
    """分析目录下的所有CSV文件"""
    results = []
    data_dict = {}  # 用于存储每个标的的数据，用于相关性分析

    # 获取所有CSV文件
    csv_files = [os.path.join(directory_path, f) for f in os.listdir(directory_path)
                 if f.lower().endswith('.csv')]

    print(f"🔍 在目录 {directory_path} 中找到 {len(csv_files)} 个CSV文件")

    # 处理每个文件
    for file_path in csv_files:
        print(f"⏳ 正在处理: {os.path.basename(file_path)}...")
        result = calculate_all_metrics(file_path)
        if result:
            results.append(result)
            # 保存数据用于相关性分析
            data_dict[result['文件名称']] = result['数据框']

    # 计算相关性
    print("\n📊 正在计算标的之间的相关性...")
    if len(data_dict) >= 2:
        corr_matrix = calculate_correlations(data_dict)
        print("相关性矩阵:")
        print(corr_matrix)
    else:
        print("⚠️ 需要至少两个有效文件才能计算相关性")
        corr_matrix = None

    return results, corr_matrix


def generate_report(results, corr_matrix=None, output_file=None):
    """生成并展示报告"""
    if not results:
        print("⚠️ 没有有效的结果可报告")
        return None

    # 创建DataFrame
    report_df = pd.DataFrame(results)

    # 添加波动率等级
    def volatility_level(vol):
        if vol is None: return 'N/A'
        if vol < 0.15:
            return '低波动'
        elif vol < 0.3:
            return '中波动'
        elif vol < 0.5:
            return '高波动'
        else:
            return '极高波动'

    report_df['波动等级'] = report_df['年化波动率'].apply(volatility_level)

    # 添加ADX趋势强度
    def adx_strength(adx):
        if adx is None: return 'N/A'
        if adx < 20:
            return '无趋势'
        elif adx < 25:
            return '弱趋势'
        elif adx < 50:
            return '强趋势'
        else:
            return '极强趋势'

    report_df['ADX强度'] = report_df['ADX(14日)'].apply(adx_strength)
    # ============ 新增：ATR筛选 ============
    # 添加高波动筛选标志
    report_df['高波动筛选'] = report_df['ATR比率(%)'].apply(
        lambda x: '是' if x is not None and x > 3.0 else '否')

    # 添加波动率比率描述
    def atr_ratio_level(ratio):
        if ratio is None: return 'N/A'
        if ratio < 1.0:
            return '极低波动'
        elif ratio < 2.0:
            return '低波动'
        elif ratio < 3.0:
            return '中波动'
        elif ratio < 5.0:
            return '高波动'
        else:
            return '极高波动'

    report_df['ATR波动等级'] = report_df['ATR比率(%)'].apply(atr_ratio_level)
    # =======================================
    # 排序
    report_df = report_df.sort_values('年化波动率', ascending=False)

    # 格式化输出
    report_df['年化波动率(%)'] = report_df['年化波动率'].apply(lambda x: f"{x * 100:.2f}%" if x is not None else 'N/A')
    report_df['ADX(14日)'] = report_df['ADX(14日)'].apply(lambda x: f"{x:.2f}" if x is not None else 'N/A')
    report_df['ATR比率(%)'] = report_df['ATR比率(%)'].apply(lambda x: f"{x:.2f}%" if x is not None else 'N/A')
    report_df['20日ATR'] = report_df['20日ATR'].apply(lambda x: f"{x:.4f}" if x is not None else 'N/A')
    report_df['现货价'] = report_df['现货价'].apply(lambda x: f"{x:.4f}" if x is not None else 'N/A')
    # 打印报告
    print("\n" + "=" * 80)
    print(f"📊 综合指标分析报告 (共 {len(report_df)} 个文件)")
    print("=" * 80)
    # 调整列顺序，包含ATR相关指标
    print(report_df[['文件名称', '数据量', '开始日期', '结束日期', '现货价',
                     '20日ATR', 'ATR比率(%)', 'ATR波动等级', '高波动筛选',
                     '年化波动率(%)', '波动等级', 'ADX(14日)', 'ADX强度',
                     '连续下跌月数']].to_string(index=False))

    # 统计摘要
    print("\n📈 波动率分布统计:")
    vol_stats = report_df['年化波动率'].describe(percentiles=[0.25, 0.5, 0.75])
    print(f"  最大值: {vol_stats['max'] * 100:.2f}%" if 'max' in vol_stats else "  最大值: N/A")
    print(f"  最小值: {vol_stats['min'] * 100:.2f}%" if 'min' in vol_stats else "  最小值: N/A")
    print(f"  平均值: {vol_stats['mean'] * 100:.2f}%" if 'mean' in vol_stats else "  平均值: N/A")
    print(f"  中位数: {vol_stats['50%'] * 100:.2f}%" if '50%' in vol_stats else "  中位数: N/A")
    # 高波动筛选结果
    high_vol_df = report_df[report_df['高波动筛选'] == '是']
    print(f"\n🔥 高波动标的 (ATR比率>3%): {len(high_vol_df)}个")
    if not high_vol_df.empty:
        print(high_vol_df[['文件名称', 'ATR比率(%)', '现货价', '20日ATR']].to_string(index=False))
    # 保存结果
    if output_file:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        final_output = f"{output_file}_{timestamp}.csv"

        # 保存主要结果
        report_df.drop(columns=['数据框'], errors='ignore').to_csv(final_output, index=False, encoding='utf-8-sig')

        # 保存相关性矩阵
        if corr_matrix is not None:
            corr_output = f"{output_file}_correlation_{timestamp}.csv"
            corr_matrix.to_csv(corr_output, encoding='utf-8-sig')
            print(f"💾 相关性矩阵已保存到: {corr_output}")

        print(f"💾 分析结果已保存到: {final_output}")
        # 保存高波动标的
        if not high_vol_df.empty:
            high_vol_output = f"{output_file}_high_volatility_{timestamp}.csv"
            high_vol_df.to_csv(high_vol_output, index=False, encoding='utf-8-sig')
            print(f"💾💾 高波动标的已保存到: {high_vol_output}")

            # 可视化高波动标的
            plt.figure(figsize=(12, 8))
            high_vol_df = high_vol_df.sort_values('ATR比率(%)', ascending=False)
            plt.bar(high_vol_df['文件名称'], high_vol_df['ATR比率(%)'])
            plt.axhline(y=3, color='r', linestyle='--', label='3%阈值')
            plt.title('高波动标的ATR比率')
            plt.ylabel('ATR/现货价 (%)')
            plt.xticks(rotation=45)
            plt.legend()
            plt.tight_layout()
            plt.savefig(f"{output_file}_high_volatility_{timestamp}.png")
            print(f"💾💾 高波动标的图表已保存到: {output_file}_high_volatility_{timestamp}.png")

    return report_df, corr_matrix
def download_stock_data(ticker,start_data,end_data):
    #下载指定股票
    stock_data = yf.download(tickers=ticker,start=start_data,end=end_data)

    #选择感兴趣的列
    data = stock_data[['Open','High','Low','Close']]
    return data
def save_to_csv(data,file_name):
    data.to_csv(file_name)


if __name__ == "__main__":
    # 设置目标目录
    target_directory = r'D:\工具\自动交易\网格交易'
    output_file = "量化分析报告"

    #抓取数据
    end_date = datetime.now()
    start_date = end_date - timedelta(days=3 * 365) #回溯3年
        #提供一个标的代码
    ticker = '512170.SH'
    #stock_data = download_stock_data(ticker,start_date,end_date)
    file_path = os.path.join(target_directory,f"{ticker}.csv")
    #save_to_csv(stock_data,file_path)

    # 执行分析
    results, corr_matrix = analyze_directory(target_directory)

    # 生成报告
    report, corr_matrix = generate_report(results, corr_matrix, output_file)

    # 示例: 如何访问特定指标
    if report is not None and not report.empty:
        sample_file = report.iloc[0]['文件名称']
        print(f"\n示例: 文件 '{sample_file}' 的分析结果:")
        print(f"  年化波动率: {report.iloc[0]['年化波动率']:.4f}")
        print(f"  ADX(14日): {report.iloc[0]['ADX(14日)']}")
        print(f"  连续下跌月数: {report.iloc[0]['连续下跌月数']}")
        print(f"  20日ATR: {report.iloc[0]['20日ATR']}")
        print(f"  现货价: {report.iloc[0]['现货价']}")
        print(f"  ATR比率: {report.iloc[0]['ATR比率(%)']}%")